#include <iostream>
#include <cstring>
#include "cmvvision.h"

/**
 * @brief 设置域值
 *
 * @param arr 域值数组
 * @param len 域值数组的长度
 * @param l 最小下标(left)
 * @param r 最大下标(right)
 * @param k 要设置的值
 */
static void set_bits(unsigned *arr, int len, int l, int r, unsigned k)
{
	l = std::max(l, 0);
	r = std::min(r + 1, len);
	for (int i = l; i < r; ++i) {
		arr[i] |= k;
	}
}

/**
 * @brief 最高有效位
 */
inline int top_bit(int n)
{
	int i = 1;

	if (!n) {
		return 0;
	}
	while (n >> i) {
		++i;
	}
	return i;
}

/**
 * @brief 最低有效位
 */
inline int bottom_bit(unsigned n)
{
	int log2modp[37] = {
		0, 1, 2,27, 3,24,28, 0, 4,17,25,31,29,12, 0,14, 5, 8,18,
		0,26,23,32,16,30,11,13, 7, 0,22,15,10, 6,21, 9,20,19
	};
	return(log2modp[(n & -n) % 37]);
}

/* 和上面的函数功能一样，直观但效率不高
inline int bottom_bit(unsigned n)
{
	int i = 0;
	if(!n) return(0);
	while(!(n&(1<<i))) i++;
	return(i + 1);
}
*/

/**
 * @brief [x, x+w)之间所有数之和
 */
inline int range_sum(int x, int w)
{
	return w * (2 * x + w - 1) / 2;
}

namespace cmv {

/**
 * @brief 初始化
 */
void Vision::initialize(const std::vector<Color> &colors, int width, int height)
{
	int size;

	m_colors = colors;
	m_width = width;
	m_height = height;
	size = m_width * m_height;

	initThreshold();

	// map
	if (m_map) {
		delete []m_map;
	}
	m_map = new unsigned[size + 1]; // 参见encoderuns()中的save变量

	// pixels
	if (m_pixels) {
		delete []m_pixels;
	}
	m_pixels = new Pixel[size];

	// runs
	if (m_runs) {
		delete []m_runs;
	}
	m_runs = new Run[size];
}

/**
 * @brief 处理一帧图像
 */
void Vision::processFrame(const unsigned char *image)
{
	int max_area;

	classifyFrame(image);
	encodeRuns();
	connectComponects();
	extractRegions();
	max_area = separateRegions();
	sortRegions(max_area);
	mergeRegions();
}

/**
 * @brief 测试阀值，输出一张rgb格式的图像，显示符合阀值的像素点
 *
 * @param image 输入图像，uyvy格式
 * @param data 输出图像，rgb格式，大小为width * height * 3
 */
void Vision::testClassify(const unsigned char *image, unsigned char *data)
{
	int size, i = 0;
	Color *color;

	classifyFrame(image);

	size = m_width * m_height;
	while (i < size) {
		while (i < size && !m_map[i]) {
			data[i * 3] = 0;
			data[i * 3 + 1] = 0;
			data[i * 3 + 2] = 0;
			++i;
		}
		while (i < size && m_map[i]) {
			color = &(m_colors[bottom_bit(m_map[i]) - 1]);
			data[i * 3] = color->color_b;
			data[i * 3 + 1] = color->color_g;
			data[i * 3 + 2] = color->color_r;
			++i;
		}
	}
}

/**
 * @brief 测试cmv结果，输出一张rgb格式的图像，显示每种颜色对应的区域
 *
 * @param image 输入图像，uyvy格式
 * @param data 输出图像，rgb格式，大小为width * height * 3
 * @param points 三维坐标点集
 */
void Vision::testProcess(const unsigned char *image, unsigned char *data, const float *depth, std::vector<Point> *points)
{
	Region *list;
	Color *color;
	int x1, y1, x2, y2;
	int step1, offset1;
	int step2, offset2;
	Point point;

	processFrame(image);

	// 背景设为黑色
	std::memset(data, 0, m_width * m_height * 3);

	// 清除points
	if (depth) {
		points->clear();
	}

	for (int i = 0; i < MAX_COLORS; ++i) { // 迭代所有颜色
		list = m_regionList[i];
		while (list) { // 迭代某种颜色的列表
			x1 = list->x1;
			y1 = list->y1;
			x2 = list->x2;
			y2 = list->y2;
			// 对于区域中的所以像素点，画出对应颜色
			for (int y = y1; y < y2; ++y) {
				step1 = m_width * y;
				step2 = step1 * 3;
				for (int x = x1; x < x2; ++x) {
					offset1 = step1 + x;
					offset2 = step2 + 3 * x;
					// 如果有符合多种颜色，会被后面的颜色覆盖
					if (m_map[offset1] & (1 << i)) {
						color = &(m_colors[i]);
						data[offset2++] = color->color_b;
						data[offset2++] = color->color_g;
						data[offset2] = color->color_r;
						if (depth && color->name == "yellow") {
							point.z = depth[offset2--];
							point.y = depth[offset2--];
							point.x = depth[offset2];
							points->push_back(point);
						}
					}
				}
			}
			// 列表中的下一个区域
			list = list->next;
		}
	}
}

/**
 * @brief 初始化域值信息
 */
void Vision::initThreshold()
{
	int size;
	unsigned k;
	Color *color;

	std::memset(m_yClass, 0, sizeof(m_yClass));
	std::memset(m_uClass, 0, sizeof(m_uClass));
	std::memset(m_vClass, 0, sizeof(m_vClass));

	size = std::min(MAX_COLORS, (int)m_colors.size());
	for (int i = 0; i < size; ++i) {
		k = (1 << i);
		color = &(m_colors[i]);
		set_bits(m_yClass, COLOR_LEVELS, color->y_min, color->y_max, k);
		set_bits(m_uClass, COLOR_LEVELS, color->u_min, color->u_max, k);
		set_bits(m_vClass, COLOR_LEVELS, color->v_min, color->v_max, k);
	}
}

/**
 * @brief 计算每个像素是否符合每种颜色的域值
 */
void Vision::classifyFrame(const unsigned char *image)
{
	int size;
	const unsigned char *addr;
	int u, y1, v, y2;
	unsigned tmp;

	size = m_width * m_height;

	for (int i = 0; i < size; i += 2) {
		addr = image + i * 2; 
		u = *addr;
		++addr;
		y1 = *addr;
		++addr;
		v = *addr;
		++addr;
		y2 = *addr;

		tmp = m_uClass[u] & m_vClass[v];
		m_map[i] = tmp & m_yClass[y1];
		m_map[i + 1] = tmp & m_yClass[y2];
	}
}

/**
 * @brief 计算每个像素是否符合每种颜色的域值，返回颜色color对应的所有像素点的坐标
 *
 * @param image yuv422格式的图像数据
 * @param color 颜色的序号（从0开始）
 * @param pixels 输出值，对应color的所有像素点的坐标
 * @param numPixels 输出值，对应pixels的个数
 */
void Vision::classifyFrame(const unsigned char *image, unsigned color, const Pixel *&pixels, int &numPixels)
{
	const unsigned char *addr = image;
	int u, y1, v, y2;
	unsigned color_bin;
	unsigned map1, map2;
	unsigned tmp;
	
	color_bin = 1 << color;
	numPixels = 0;

	for (int i = 0; i < m_height; ++i) {
		// 不考虑m_width为奇数的情况
		for (int j = 0; j < m_width; j += 2) {
			u = *addr;
			++addr;
			y1 = *addr;
			++addr;
			v = *addr;
			++addr;
			y2 = *addr;
			++addr;

			tmp = m_uClass[u] & m_vClass[v];

			map1 = tmp & m_yClass[y1];
			if (map1 & color_bin) {
				m_pixels[numPixels].x = j;
				m_pixels[numPixels].y = i;
				++numPixels;
			}
			m_map[i] = map1;

			map2 = tmp & m_yClass[y2];
			if (map2 & color_bin) {
				m_pixels[numPixels].x = j + 1;
				m_pixels[numPixels].y = i;
				++numPixels;
			}
			m_map[i + 1] = map2; 
		}
	}

	pixels = m_pixels;
}

/**
 * @brief 计算规定范围内每个像素是否符合每种颜色的域值，返回颜色color对应的所有像素点的坐标
 *
 * @param image yuv422格式的图像数据
 * @param robot_pix_x 机器人在全景图像中的x坐标
 * @param robot_pix_y 机器人在全景图像中的y坐标
 * @param r_big 大正方形内切圆的半径，如果为0表示没有外边框
 * @param r_small 小正方形内切圆的半径（这里没有对r_big和r_small进行检查，调用时必须保证其合理性）
 * @param color 颜色的序号（从0开始）
 * @param pixels 输出值，对应矩形框和color的所有像素点的坐标
 * @param numPixels 输出值，对应pixels的个数
 */
void Vision::classifyFrame(const unsigned char *image, int robot_pix_x, int robot_pix_y, int r_big, int r_small, unsigned color,
	const Pixel *&pixels, int &numPixels)
{
	int x1_big, y1_big, x2_big, y2_big; // 大正方形对角点的坐标
	int x1_small, y1_small, x2_small, y2_small; // 小正方形对角点的坐标
	const unsigned char *addr1, *addr2;
	int yuv_u, yuv_y1, yuv_v, yuv_y2;
	unsigned map1, map2;
	unsigned color_bin;
	unsigned tmp;

	// 计算正方形对角点坐标
	if (r_big == 0) {
		x1_big = 0;
		y1_big = 0;
		x2_big = m_width;
		y2_big = m_height;
	} else {
		x1_big = robot_pix_x - r_big;
		y1_big = robot_pix_y - r_big;
		x2_big = robot_pix_x + r_big;
		y2_big = robot_pix_y + r_big;
	}
	x1_small = robot_pix_x - r_small;
	y1_small = robot_pix_y - r_small;
	x2_small = robot_pix_x + r_small;
	y2_small = robot_pix_y + r_small;

	// 初始化
	color_bin = 1 << color;
	numPixels = 0;

	// 两个正方形之间的区域可以划分为四个部分，分别进行计算

	// 第一个矩形
	for (int i = y1_big; i < y1_small; ++i) {
		addr1 = image + m_width * i * 2;
		for (int j = x1_big; j < x2_big; j += 2) { // 可能最后一列不会扫描
			addr2 = addr1 + j * 2;
			yuv_u = *addr2;
			++addr2;
			yuv_y1 = *addr2;
			++addr2;
			yuv_v = *addr2;
			++addr2;
			yuv_y2 = *addr2;

			tmp = m_uClass[yuv_u] & m_vClass[yuv_v];
			map1 = tmp & m_yClass[yuv_y1];
			map2 = tmp & m_yClass[yuv_y2];
			if (map1 & color_bin) {
				m_pixels[numPixels].x = j;
				m_pixels[numPixels].y = i;
				++numPixels;
			}
			if (map2 & color_bin) {
				m_pixels[numPixels].x = j + 1;
				m_pixels[numPixels].y = i;
				++numPixels;
			}
		}
	}

	// 第二个矩形
	for (int i = y1_small; i < y2_small; ++i) {
		addr1 = image + m_width * i * 2;
		for (int j = x1_big; j < x1_small; j += 2) { // 可能最后一列不会扫描
			addr2 = addr1 + j * 2;
			yuv_u = *addr2;
			++addr2;
			yuv_y1 = *addr2;
			++addr2;
			yuv_v = *addr2;
			++addr2;
			yuv_y2 = *addr2;

			tmp = m_uClass[yuv_u] & m_vClass[yuv_v];
			map1 = tmp & m_yClass[yuv_y1];
			map2 = tmp & m_yClass[yuv_y2];
			if (map1 & color_bin) {
				m_pixels[numPixels].x = j;
				m_pixels[numPixels].y = i;
				++numPixels;
			}
			if (map2 & color_bin) {
				m_pixels[numPixels].x = j + 1;
				m_pixels[numPixels].y = i;
				++numPixels;
			}
		}
	}

	// 第三个矩形
	for (int i = y1_small; i < y2_small; ++i) {
		addr1 = image + m_width * i * 2;
		for (int j = x2_small; j < x2_big; j += 2) { // 可能最后一列不会扫描
			addr2 = addr1 + j * 2;
			yuv_u = *addr2;
			++addr2;
			yuv_y1 = *addr2;
			++addr2;
			yuv_v = *addr2;
			++addr2;
			yuv_y2 = *addr2;

			tmp = m_uClass[yuv_u] & m_vClass[yuv_v];
			map1 = tmp & m_yClass[yuv_y1];
			map2 = tmp & m_yClass[yuv_y2];
			if (map1 & color_bin) {
				m_pixels[numPixels].x = j;
				m_pixels[numPixels].y = i;
				++numPixels;
			}
			if (map2 & color_bin) {
				m_pixels[numPixels].x = j + 1;
				m_pixels[numPixels].y = i;
				++numPixels;
			}
		}
	}

	// 第四个矩形
	for (int i = y2_small; i < y2_big; ++i) {
		addr1 = image + m_width * i * 2;
		for (int j = x1_big; j < x2_big; j += 2) { // 可能最后一列不会扫描
			addr2 = addr1 + j * 2;
			yuv_u = *addr2;
			++addr2;
			yuv_y1 = *addr2;
			++addr2;
			yuv_v = *addr2;
			++addr2;
			yuv_y2 = *addr2;

			tmp = m_uClass[yuv_u] & m_vClass[yuv_v];
			map1 = tmp & m_yClass[yuv_y1];
			map2 = tmp & m_yClass[yuv_y2];
			if (map1 & color_bin) {
				m_pixels[numPixels].x = j;
				m_pixels[numPixels].y = i;
				++numPixels;
			}
			if (map2 & color_bin) {
				m_pixels[numPixels].x = j + 1;
				m_pixels[numPixels].y = i;
				++numPixels;
			}
		}
	}

	pixels = m_pixels;
}

/**
 * @brief 逐行扫描得到runs
 */
void Vision::encodeRuns()
{
	unsigned *row; // 数组行头的地址
	int start; // run的开始下标
	int count = 0; // run的数目
	unsigned value, save;

	save = m_map[0];

	for (int i = 0; i < m_height; ++i) { // 每行
		row = &m_map[i * m_width];
		// save保存下一行行首的值
		// 用来结束一行数据的扫描
		row[0] = save;
		save = row[m_width];
		row[m_width] = (unsigned)-1;

		for (int j = 0; j < m_width;) { // 每列
			value = row[j];
			start = j;
			++j;
			while (row[j] == value) {
				++j;
			}

			m_runs[count].color = value;
			m_runs[count].length = j - start;
			m_runs[count].parent = count;
			++count;
		}
	}

	m_numRuns = count;
}

/**
 * @brief 区域连接标记(connected component labelling)
 */
void Vision::connectComponects()
{
	int start1, start2; // run的起始位置
	int index1, index2, index; // m_runs的下标，index1的行号大于或等于index2的行号
	int parent1, parent2, parent;
	Run run1, run2;

	start1 = start2 = 0;
	index1 = index2 = 0;

	// 第一行
	while (start1 < m_width) {
		start1 += m_runs[index1++].length;
	}
	start1 = 0;

	// 剩下的行
	run1 = m_runs[index1]; // 第二行的第一个run
	run2 = m_runs[index2]; // 第一行的第一个run
	index = index1;
	while (index1 < m_numRuns) {
		// 更新parent
		if (run1.color && run1.color == run2.color // 相同颜色归属的run
				// 相邻的run
				&& ((start1 >= start2 && start1 < start2 + run2.length)
					|| (start2 >= start1 && start2 < start1 + run1.length))) {
			if (index != index1) { // 从未扫描过的run
				// 未扫描过的run的parent肯定比扫描过的要大（初始化时设定的）
				m_runs[index1].parent = run1.parent = run2.parent;
				index = index1;
			} else { // 已经扫描过的run
				parent1 = run1.parent;
				while (parent1 != m_runs[parent1].parent) {
					parent1 = m_runs[parent1].parent;
				}
				parent2 = run2.parent;
				while (parent2 != m_runs[parent2].parent) {
					parent2 = m_runs[parent2].parent;
				}
				// 设置为较小的parent值
				if (parent1 < parent2) {
					m_runs[parent2].parent = parent1;
				} else {
					m_runs[parent1].parent = parent2;
				}
			}
		}

		// 转到run1的下一个run或者run2的下一个run
		if (start1 + run1.length < start2 + run2.length) {
			start1 += run1.length;
			run1 = m_runs[++index1];
		} else {
			start2 += run2.length;
			run2 = m_runs[++index2];
		}
	}

	// 路径压缩
	for (int i = 0; i < m_numRuns; ++i) {
		parent = m_runs[i].parent;
		if (parent > i) { // 有这种可能性么？？？
			std::cout << "what!!!!" << std::endl;
			while (parent != m_runs[parent].parent) {
				parent = m_runs[parent].parent;
			}
			m_runs[i].parent = parent;
		} else {
			m_runs[i].parent = m_runs[parent].parent;
		}
	}

	// Ouch, my brain hurts too. (^_^)
}

/**
 * @brief 计算连通区域
 */
void Vision::extractRegions()
{
	Run run;
	int id;
	int count = 0;
	int x = 0, y = 0;
	int tmp;

	// 扫描每个run
	for (int i = 0; i < m_numRuns; ++i) {
		run = m_runs[i];
		if (run.color) {
			if (run.parent == i) { // 父亲是本身，添加一个新的连通区域
				m_runs[i].parent = id = count; // 重新标号run
				m_regions[id].color = bottom_bit(run.color) - 1;
				m_regions[id].area = run.length;
				m_regions[id].x1 = x;
				m_regions[id].y1 = y;
				m_regions[id].x2 = x + run.length;
				m_regions[id].y2 = y;
				m_regions[id].sum_x = range_sum(x, run.length);
				m_regions[id].sum_y = y * run.length;
				++count;
				if (count >= MAX_REGIONS) {
					m_numRegions = MAX_REGIONS;
					return;
				}
			} else { // 父亲不是本身，添加run到已有的连通区域
				id = m_runs[run.parent].parent;
				m_runs[i].parent = id; // 重新标号run
				m_regions[id].area += run.length;
				m_regions[id].x2 = std::max(x + run.length, m_regions[id].x2);
				m_regions[id].x1 = std::min(x, m_regions[id].x1);
				m_regions[id].y2 = y;
				m_regions[id].sum_x += range_sum(x, run.length);
				m_regions[id].sum_y += y * run.length;
			}
		}

		x = (x + run.length) % m_width;
		y += (x == 0);
	}

	// 计算中心点
	for (int i = 0; i < count; ++i) {
		tmp = m_regions[i].area;
		m_regions[i].cen_x = (float)m_regions[i].sum_x / tmp;
		m_regions[i].cen_y = (float)m_regions[i].sum_y / tmp;
	}

	m_numRegions = count;
}

/**
 * @brief 将区域按照颜色进行分类，每种颜色对应一个区域列表
 */
int Vision::separateRegions()
{
	Region *region;
	int area, max_area = 0;
	int color_id;

	// 初始化
	for (int i = 0; i < MAX_COLORS; ++i) {
		m_regionCount[i] = 0;
		m_regionList[i] = NULL;
	}

	// 扫描每一个区域
	for (int i = 0; i < m_numRegions; ++i) {
		region = &m_regions[i];
		area = region->area;
		if (area >= MIN_AREA_SIZE) {
			if (area > max_area) {
				max_area = area;
			}
			// 添加一个区域到对应的颜色分类
			color_id = region->color;
			++m_regionCount[color_id];
			region->next = m_regionList[color_id];
			m_regionList[color_id] = region;
		}
	}

	return max_area;
}

/**
 * @brief 对每种颜色对应的区域列表按照尺寸从大到小排序
 */
void Vision::sortRegions(int max_area)
{
	int passes;

	passes = top_bit((max_area + RBITS - 1) / RBITS);
	for (int i = 0; i < MAX_COLORS; ++i) {
		m_regionList[i] = sortRegionListByArea(m_regionList[i], passes);
	}
}

/**
 * @brief 合并相同颜色的区域
 */
void Vision::mergeRegions()
{
	int count;

	for (int i = 0; i < MAX_COLORS; ++i) {
		count = mergeRegions(m_regionList[i], 1000000, 0.5);
		m_regionCount[i] -= count;
	}
}

/**
 * @brief 对一个区域列表进行基数排序（对应一种颜色）
 */
Region* Vision::sortRegionListByArea(Region *list, int passes)
{
	Region *table[RADIX], *region, *region_next;
	int slot, shift;

	if (!list || !list->next) {
		return list;
	}

	// 初始化
	for (int i = 0; i < RADIX; ++i) {
		table[i] = NULL;
	}

	for (int i = 0; i < passes; ++i) {
		shift = RBITS * i; // 当前的pass对应的shift
		region = list;
		while (region) {
			region_next = region->next;
			slot = ((region->area) >> shift) & RMASK; // 相当于% RADIX
			// 将region添加到slot槽中
			region->next = table[slot];
			table[slot] = region;
			// 列表中的下一个区域
			region = region_next;
		}

		// 使用已进行过的pass排序
		list = NULL;
		for (int j = 0; j < RADIX; ++j) {
			region = table[j];
			table[j] = NULL;
			while (region) {
				region_next = region->next;
				// 添加region到list中
				region->next = list;
				list = region;
				// 下一个region
				region = region_next;
			}
		}
	}

	return list;
}

/**
 * @brief 合并相邻的相同颜色的连通区域
 *
 * @param p 同一个颜色的连通区域列表
 * @param num 最大合并次数
 * @param density_thresh 密度阀值（像素点个数除以区域边框的面积）
 *
 * @return 合并次数，等于被移除的区域的个数
 */
int Vision::mergeRegions(Region *p, int num, double density_thresh)
{
	Region *q, *s;
	int l, r, t, b; // 合并之后的边框
	int a; // 面积
	int merged; // 合并次数

	merged = 0;

	// 算法复杂度n^2
	while (p && merged < num) {
		q = p->next;
		s = p;

		while (q) { // 将p与列表中剩下的区域进行比较和合并
			l = std::min(p->x1, q->x1);
			r = std::max(p->x2, q->x2);
			t = std::min(p->y1, q->y1);
			b = std::max(p->y2, q->y2);
			a = (r - l) * (b - t);
			
			if ((double)(p->area + q->area) / a > density_thresh) {
				// 密度大于阀值，合并
				a = p->area + q->area;
				p->x1 = l;
				p->x2 = r;
				p->y1 = t;
				p->y2 = b;
				p->cen_x = ((p->cen_x * p->area) + (q->cen_x * q->area)) / a;
				p->cen_y = ((p->cen_y * p->area) + (q->cen_y * q->area)) / a;
				p->area = a;
				// 移除q
				q = q->next;
				s->next = q; // 从list中移除
				// 合并次数加1
				++merged;
			} else {
				// 不合并
				s = q;
				q = q->next;
			}
		}

		// 迭代p
		p = p->next;
	}

	return merged;
}

}
